package cg.cmm.rag.controller;

import cg.cmm.rag.dto.Person;
import cg.cmm.rag.service.Customer;
import cg.cmm.rag.service.PersonService;
import cg.cmm.rag.service.Writer;
import cg.cmm.rag.tool.ToolUtils;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentBySentenceSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.image.Image;
import dev.langchain4j.data.message.*;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.request.ResponseFormatType;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.chat.request.json.JsonSchema;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.image.ImageModel;
import dev.langchain4j.model.moderation.Moderation;
import dev.langchain4j.model.moderation.ModerationModel;
import dev.langchain4j.model.openai.*;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.web.search.WebSearchTool;
import dev.langchain4j.web.search.searchapi.SearchApiWebSearchEngine;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.lang.reflect.InvocationTargetException;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author: baibing
 * create time: 2025-03-07 23:10
 * description:
 */
@RestController
public class ChatController {

    ChatLanguageModel chatLanguageModel;
    StreamingChatLanguageModel streamingChatLanguageModel;
    ModerationModel moderationModel;
    ImageModel imageModel;
    ChatMemory chatMemory;
    ChatMemory tokenWindowChatMemory;

    private final Map<String, ChatMemory> chatMemoryMap = new ConcurrentHashMap<>();

    public ChatController() {
        this.chatLanguageModel = OpenAiChatModel.builder()
                .timeout(Duration.ofSeconds(360))
//                .baseUrl("https://api.deepseek.com/v1")
//                .apiKey("sk-304d5969ac1c4ebe8ebf26f8d8309536")
//                .modelName("deepseek-chat")
                .baseUrl("http://localhost:11434/v1") // Ollama 本地服务地址
                .modelName("qwen2.5:latest") // 模型名称
                .apiKey("your-api-key") // 如果需要 API 密钥
                .build();

        this.streamingChatLanguageModel = OpenAiStreamingChatModel.builder()
                .timeout(Duration.ofSeconds(360))
                .baseUrl("http://localhost:11434/v1") // Ollama 本地服务地址
                .modelName("qwen2.5:latest") // 模型名称
                .apiKey("your-api-key") // 如果需要 API 密钥
                .build();
        this.moderationModel = OpenAiModerationModel.builder()
                .timeout(Duration.ofSeconds(360))
                .baseUrl("http://localhost:11434/v1") // Ollama 本地服务地址
                .modelName("qwen2.5:latest") // 模型名称
                .apiKey("your-api-key") // 如果需要 API 密钥
                .build();
        this.imageModel = OpenAiImageModel.builder()
                .timeout(Duration.ofSeconds(360))
                .baseUrl("http://localhost:11434/v1") // Ollama 本地服务地址
                .modelName("qwen2.5:latest") // 模型名称
                .apiKey("your-api-key") // 如果需要 API 密钥
                .build();
        this.chatMemory = MessageWindowChatMemory.builder()
                .maxMessages(10)
                .build();

        this.tokenWindowChatMemory = TokenWindowChatMemory.builder()
                .maxTokens(1000, new OpenAiTokenizer())
                .build();

    }

    /**
     * 测试对话
     *
     * @param message
     * @return
     */

    @GetMapping("/message")
    public String model(@RequestParam(value = "message", defaultValue = "Hello") String message) {
        SystemMessage systemMessage = SystemMessage.from("你是一名java架构师");
        UserMessage userMessage = UserMessage.from(message);
        ChatResponse response = chatLanguageModel.chat(systemMessage, userMessage);
        return response.aiMessage().text();
    }

    /**
     * 测试流式对话
     *
     * @param message
     * @return
     */
    @GetMapping("/stream")
    public String Stringream(@RequestParam(value = "message", defaultValue = "Hello") String message) {
        SystemMessage systemMessage = SystemMessage.from("你是一名java架构师");
        UserMessage userMessage = UserMessage.from(message);
        streamingChatLanguageModel.chat(List.of(systemMessage, userMessage), new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String partialResponse) {
                System.out.println(partialResponse);
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                System.out.println(completeResponse.aiMessage().text());
            }

            @Override
            public void onError(Throwable error) {

            }
        });
        return "stream";
    }

    /**
     * 去掉敏感词提问
     *
     * @param message
     * @return
     */
    @GetMapping("/moderation")
    public String moderation(@RequestParam(value = "message", defaultValue = "Hello") String message) {
        SystemMessage systemMessage = SystemMessage.from("你是一名java架构师");
        UserMessage userMessage = UserMessage.from(message);
        Response<Moderation> response = moderationModel.moderate(List.of(systemMessage, userMessage));
        return response.content().flaggedText();
    }

    /**
     * @param message
     * @return
     */
    @GetMapping("/image")
    public String image(@RequestParam(value = "message", defaultValue = "Hello") String message) {
        Response<Image> response = imageModel.generate("一只功夫熊猫");
        return response.content().url().toString();
    }

    /**
     * 测试聊天记忆
     *
     * @param message
     * @return
     */
    @GetMapping("/memory")
    public String memory(@RequestParam(value = "message", defaultValue = "Hello") String message) {
        if (chatMemory.messages().isEmpty()) {
            SystemMessage systemMessage = SystemMessage.from("你是一名java架构师");
            chatMemory.add(systemMessage);
        }
        UserMessage userMessage = UserMessage.from(message);
        chatMemory.add(userMessage);
        ChatResponse response = chatLanguageModel.chat(chatMemory.messages());
        chatMemory.add(response.aiMessage());

        StringBuilder sb = new StringBuilder();
        for (ChatMessage chatMessage : chatMemory.messages()) {
            sb.append(chatMessage.toString()).append("\n");
        }
        return sb.toString();
    }

    @GetMapping("/token_memory")
    public String tokenMemory(@RequestParam(value = "message", defaultValue = "Hello") String message) {
        if (tokenWindowChatMemory.messages().isEmpty()) {
            SystemMessage systemMessage = SystemMessage.from("你是一名java架构师");
            tokenWindowChatMemory.add(systemMessage);
        }
        UserMessage userMessage = UserMessage.from(message);
        tokenWindowChatMemory.add(userMessage);
        ChatResponse response = chatLanguageModel.chat(tokenWindowChatMemory.messages());
        tokenWindowChatMemory.add(response.aiMessage());
        StringBuilder sb = new StringBuilder();
        for (ChatMessage chatMessage : tokenWindowChatMemory.messages()) {
            sb.append(chatMessage.text()).append("\n");
        }
        return sb.toString();
    }

    /**
     * 测试聊天记忆
     *
     * @param message
     * @return
     */
    @GetMapping("/memory_id")
    public String memoryId(@RequestParam(value = "message", defaultValue = "Hello") String message, @RequestParam(value = "memoryId") String memoryId) {
        MessageWindowChatMemory idMemory = MessageWindowChatMemory.builder()
                .maxMessages(10)
                .id(memoryId)
                .build();
        if (idMemory.messages().isEmpty()) {
            SystemMessage systemMessage = SystemMessage.from("你是一名java架构师");
            idMemory.add(systemMessage);
        }
        UserMessage userMessage = UserMessage.from(message);
        idMemory.add(userMessage);
        ChatResponse response = chatLanguageModel.chat(idMemory.messages());
        idMemory.add(response.aiMessage());
        StringBuilder sb = new StringBuilder();
        for (ChatMessage chatMessage : idMemory.messages()) {
            sb.append(chatMessage.text()).append("\n");
        }
        return sb.toString();
    }

    @GetMapping("/tool")
    public String tool(@RequestParam(value = "message", defaultValue = "Hello") String message) {
        List<ToolSpecification> toolSpecifications = null;
        try {
            toolSpecifications = ToolSpecifications.toolSpecificationsFrom(ToolUtils.class);
        } catch (Exception e) {
            e.printStackTrace();
        }
        if (chatMemory.messages().isEmpty()) {
            SystemMessage systemMessage = SystemMessage.from("你是一名java架构师");
            chatMemory.add(systemMessage);
        }
        UserMessage userMessage = UserMessage.from(message);
        chatMemory.add(userMessage);
        Response<AiMessage> response = chatLanguageModel.generate(chatMemory.messages(), toolSpecifications);
        if (response.content().hasToolExecutionRequests()) {
            for (ToolExecutionRequest toolExecutionRequest : response.content().toolExecutionRequests()) {
                try {
                    String result = (String) ToolUtils.class.getMethod(toolExecutionRequest.name()).invoke(null);
                    ToolExecutionResultMessage toolExecutionResultMessage = ToolExecutionResultMessage.from(toolExecutionRequest.id(), toolExecutionRequest.name(), result);
                    Response<AiMessage> response1 = chatLanguageModel.generate(List.of(userMessage, response.content(), toolExecutionResultMessage));
                    chatMemory.add(response1.content());
                } catch (IllegalAccessException e) {
                    e.printStackTrace();
                } catch (InvocationTargetException e) {
                    e.printStackTrace();
                } catch (NoSuchMethodException e) {
                    e.printStackTrace();
                }
            }
        } else {
            chatMemory.add(response.content());
        }
        StringBuilder sb = new StringBuilder();
        for (ChatMessage chatMessage : chatMemory.messages()) {
            sb.append(chatMessage.toString()).append("\n");
        }
        return sb.toString();
    }

    @GetMapping("/tool2")
    public String tool2(@RequestParam(value = "message", defaultValue = "Hello") String message) {
        List<ToolSpecification> toolSpecifications = null;
        try {
            toolSpecifications = ToolSpecifications.toolSpecificationsFrom(ToolUtils.class);
        } catch (Exception e) {
            e.printStackTrace();
        }
        if (chatMemory.messages().isEmpty()) {
            SystemMessage systemMessage = SystemMessage.from("你是一名java架构师");
            chatMemory.add(systemMessage);
        }
        UserMessage userMessage = UserMessage.from(message);
        chatMemory.add(userMessage);
        Response<AiMessage> response = chatLanguageModel.generate(chatMemory.messages(), toolSpecifications);
        if (response.content().hasToolExecutionRequests()) {
            for (ToolExecutionRequest toolExecutionRequest : response.content().toolExecutionRequests()) {
                try {
                    String result = (String) ToolUtils.class.getMethod(toolExecutionRequest.name()).invoke(null);
                    ToolExecutionResultMessage toolExecutionResultMessage = ToolExecutionResultMessage.from(toolExecutionRequest.id(), toolExecutionRequest.name(), result);
                    Response<AiMessage> response1 = chatLanguageModel.generate(List.of(userMessage, response.content(), toolExecutionResultMessage));
                    chatMemory.add(response1.content());
                } catch (IllegalAccessException e) {
                    e.printStackTrace();
                } catch (InvocationTargetException e) {
                    e.printStackTrace();
                } catch (NoSuchMethodException e) {
                    e.printStackTrace();
                }
            }
        } else {
            chatMemory.add(response.content());
        }
        StringBuilder sb = new StringBuilder();
        for (ChatMessage chatMessage : chatMemory.messages()) {
            sb.append(chatMessage.toString()).append("\n");
        }
        return sb.toString();
    }

    @GetMapping("/ai_service")
    public String aiService(@RequestParam(value = "message", defaultValue = "Hello") String message) {
        Writer writer = AiServices.create(Writer.class, chatLanguageModel);
        return writer.write(message, "最美的人", "1000");
    }

    @GetMapping("/ai_service2")
    public String aiService2(@RequestParam(value = "message", defaultValue = "Hello") String message, @RequestParam(value = "memoryId") String memoryId) {
//        MessageWindowChatMemory.withMaxMessages(10);
//        ChatMemory chatMemory = chatMemoryMap.computeIfAbsent(memoryId, id ->
//                MessageWindowChatMemory.builder().id(id).build());
        Customer customer = AiServices.builder(Customer.class)
                .chatLanguageModel(chatLanguageModel)
//                .chatMemory(chatMemory)
                .chatMemoryProvider(id -> {
                    // 设置一个默认的maxMessages值，例如10
                    return chatMemoryMap.computeIfAbsent(memoryId, mid ->
                            MessageWindowChatMemory.builder()
                                    .id(mid)
                                    .maxMessages(10) // 确保设置了maxMessages
                                    .build());
                })
                .tools(new ToolUtils())
                .build();
        return customer.answer(message);
    }

    /**
     * 根据用户id获取历史对话
     *
     * @param message
     * @param memoryId
     * @param systemMessage
     * @return dev.langchain4j.service.DefaultAiServices 核心这个类处理业务，跟下代码
     */
    @GetMapping("/answerByUserId")
    public String answerByUserId(@RequestParam(value = "message", defaultValue = "Hello") String message, @RequestParam(value = "memoryId") String memoryId, @RequestParam(value = "systemMessage", defaultValue = "") String systemMessage) {

        Customer customer = AiServices.builder(Customer.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemoryProvider(id -> {
                    return MessageWindowChatMemory.withMaxMessages(10);
                })
                .tools(new ToolUtils())
                .build();
        return customer.answerByUserId(memoryId, message, systemMessage);
    }

    @GetMapping("/embedding")
    public String embedding(@RequestParam(value = "message", defaultValue = "Hello") String message) {

        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .timeout(Duration.ofSeconds(360))
                .baseUrl("http://localhost:11434/v1") // Ollama 本地服务地址
                .modelName("qwen2.5:latest") // 模型名称
                .apiKey("your-api-key") // 如果需要 API 密钥
                .build();

        Response<Embedding> response = embeddingModel.embed(message);
        System.out.println(response.content().vector().length);
        return response.content().toString();
    }


    @Autowired
    public EmbeddingStore<TextSegment> embeddingStore;

    @GetMapping("/loadragdoc")
    public String loadragdoc() {
        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .timeout(Duration.ofSeconds(360))
                .baseUrl("http://localhost:11434/v1") // Ollama 本地服务地址
                .modelName("qwen2.5:latest") // 模型名称
                .apiKey("your-api-key") // 如果需要 API 密钥
                .build();
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .embeddingModel(embeddingModel)  // 确保这里的embeddingModel不是null
                .embeddingStore(embeddingStore)
                .documentSplitter(new DocumentBySentenceSplitter(1000, 300))
                .build();

        List<Document> documents = FileSystemDocumentLoader.loadDocuments("D:\\DevProject\\cgframework\\cg-cmm\\cg-cmm-langchain4j-rag\\documents");
        ingestor.ingest(documents);

        return "加载成功";
    }

    @Autowired
    public Customer customer;

    @GetMapping("/answerByUserIdRag")
    public String answerByUserIdRag(@RequestParam(value = "message", defaultValue = "Hello") String message, @RequestParam(value = "memoryId") String memoryId, @RequestParam(value = "systemMessage", defaultValue = "") String systemMessage) {
        return customer.answerByUserId(memoryId, message, systemMessage);
    }

    @GetMapping("/json")
    public String json(@RequestParam(value = "message", defaultValue = "Hello") String message, @RequestParam(value = "memoryId") String memoryId, @RequestParam(value = "systemMessage", defaultValue = "") String systemMessage) {
        ResponseFormat responseFormat = ResponseFormat.builder()
                .type(ResponseFormatType.JSON)
                .jsonSchema(JsonSchema.builder()
                        .rootElement(JsonObjectSchema.builder()
                                .addIntegerProperty("age", "年龄")
                                .addIntegerProperty("height", "身高")
                                .required("age", "height")
                                .build())
                        .build())
                .build();
        ChatResponse chat = chatLanguageModel.chat(ChatRequest.builder()
                .messages(List.of(UserMessage.from(message)))
                .parameters(ChatRequestParameters.builder().responseFormat(responseFormat).build())
                .build());
        return chat.aiMessage().text();
    }

    @GetMapping("/json1")
    public String json1(@RequestParam(value = "message", defaultValue = "Hello") String message, @RequestParam(value = "memoryId") String memoryId, @RequestParam(value = "systemMessage", defaultValue = "") String systemMessage) {
        PersonService personService = AiServices.builder(PersonService.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemoryProvider(id -> {
                    return MessageWindowChatMemory.withMaxMessages(10);
                })
                .tools(new ToolUtils())
                .build();
        Person person = personService.extracePerson(message);
        return person.toString();
    }


    @Autowired
    public SearchApiWebSearchEngine engine;
    @GetMapping("/web_search")
    public String webSearch(@RequestParam(value = "message", defaultValue = "Hello") String message, @RequestParam(value = "memoryId") String memoryId, @RequestParam(value = "systemMessage", defaultValue = "") String systemMessage) {
        Customer customer = AiServices.builder(Customer.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemoryProvider(id -> {
                    return MessageWindowChatMemory.withMaxMessages(10);
                })
                .tools(new ToolUtils(),new WebSearchTool(engine))
                .build();
        return customer.answerByUserId(memoryId, message, systemMessage);
    }
}
